import os
import matplotlib.pyplot as plt
import numpy as np


def moving_average(data, window_size=10):
    if len(data) < window_size:
        return np.array(data)
    return np.convolve(data, np.ones(window_size) / window_size, mode='valid')


def plot_single_metric(data, title, ylabel, save_path, window=15):
    plt.figure(figsize=(14, 5))
    x = np.arange(len(data))
    y = np.array(data)

    # 平滑曲线
    y_smooth = moving_average(y, window)
    x_smooth = np.arange(len(y_smooth))

    # 原始散点（透明灰色）
    plt.scatter(x, y, s=8, alpha=0.2, color='gray', label="Raw")
    plt.plot(x_smooth, y_smooth, color='tab:blue', linewidth=2, label="Smoothed")

    plt.title(title, fontsize=16)
    plt.xlabel("Episode Index", fontsize=13)
    plt.ylabel(ylabel, fontsize=13)
    plt.legend(fontsize=12)
    plt.grid(True, linestyle="--", alpha=0.4)
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def plot_metrics(metrics_dict, save_dir):
    os.makedirs(save_dir, exist_ok=True)

    if "rewards" in metrics_dict:
        plot_single_metric(
            metrics_dict["rewards"],
            "Episode Rewards",
            "Reward",
            os.path.join(save_dir, "rewards.png")
        )

    if "wins" in metrics_dict:
        plot_single_metric(
            metrics_dict["wins"],
            "Episode Win Flags",
            "Win (1) or Loss (0)",
            os.path.join(save_dir, "wins.png")
        )

    if "kd_ratios" in metrics_dict:
        plot_single_metric(
            metrics_dict["kd_ratios"],
            "Episode K/D Ratios",
            "K/D Ratio",
            os.path.join(save_dir, "kd_ratios.png")
        )

    if "entropies" in metrics_dict:
        plot_single_metric(
            metrics_dict["entropies"],
            "Policy Entropy",
            "Entropy",
            os.path.join(save_dir, "entropies.png")
        )
